Add transform_values UDF#22689
Conversation
Introduce a new higher-order UDF that returns a new map by applying a lambda `(k, v) -> expr` to each entry of the input map, transforming the values while preserving the keys. Registered under the alias `transform_values` via `with_aliases` from apache#22593. Also fix a latent indexing bug in `LambdaExpr`: when a multi-parameter lambda referenced only a subset of its declared parameters (e.g. `(k, v) -> v * 2`), the compressed column index map could shift lambda params into outer capture slots. `LambdaExpr` now tracks `outer_columns_count` and remaps indices so outer captures and lambda params keep stable positions regardless of which params are used.
Keep a single canonical name `map_transform` and remove the `transform_values` alias to simplify registration (no more `Arc::new(... .with_aliases(...))` dance in `lib.rs`).
Add map_transform UDF
Run `./dev/update_function_docs.sh` so the user-facing function reference picks up the new `map_transform` entry from its `#[user_doc(...)]` annotation.
Rename `map_transform` to `transform_values` (Spark/Trino spelling, leaves room for a future `transform_keys`), extract shared helpers in this file and `utils.rs` (`get_map_key_value_fields`, reuse `value_lambda_pair` from `lambda_utils`), return an array of nulls (instead of a null scalar) when every input row is null, and move the empty-entries fast path up next to the all-null fast path. Also add a sqllogictest covering the all-null rows behavior and regenerate the scalar function docs.
There was a problem hiding this comment.
This lgtm, from the understanding the bug this PR fixes (aside from adding transform_values) is:
- We have
map_transform(my_map, (k, v) -> v * 2) - The
Recordbatchthe lambda processes has two columns k and v (indices 0,1) LambdaExprprojects the batch so it keeps the columns the lambda body actually references (v with index 1)- v ends up with index 0 in the projected schema
- The body will read the wrong column in index 0 (which is k and not v)
LiaCastaneda
left a comment
There was a problem hiding this comment.
@gabotechs or @rluvaton -- whenever you have time, could you take a look at this PR? It adds transform_values and fixes a lambda indexing bug
| } | ||
|
|
||
| #[cfg(test)] | ||
| mod tests { |
There was a problem hiding this comment.
clear_null_values currently doesn't support MapArray (I will create a PR now to add support for that, so the test will probably fail until my pr is merged).
Please add tests like:
datafusion/datafusion/functions-nested/src/array_transform.rs
Lines 223 to 293 in 34834dc
There was a problem hiding this comment.
Created a pr to cleanup here:
Adds four unit tests requested in PR review: - `transform_values_empty_map`: single row, empty entries - `transform_values_mixed_empty_null_populated_rows`: empty + null + populated rows in one batch to exercise offset/null-mask plumbing - `transform_values_on_sliced_map_should_not_evaluate_on_unreachable_values`: slices a 4-row map and uses `100 / v` so a leading `0` would divide-by-zero if offset handling is wrong - `transform_values_function_should_not_be_evaluated_on_values_underlying_null` (ignored): documents the dependency on apache#22847 which adds `clear_null_values` support for `MapArray` While adding the empty-map test, the previously-existing "empty entries" fast-path was found to be broken — `ScalarValue::new_default` for a `Map` return type produces a scalar that decodes back as a `Struct`, not a `Map`, and the test downcast failed. The fast path was just an optimisation skipping the lambda evaluation; the regular path already handles empty inputs correctly, so the buggy fast path was removed.
Per rluvaton's nit on PR apache#22689: the all-null fast path can return a typed null scalar (`ScalarValue::try_new_null(return_type)`) and let the caller broadcast it back out to the input row count, instead of materialising a full `new_null_array` ourselves.
|
Can you please extract the breaking change parts to a different PR since it requires a special attention as oppose to adding UDF that does not affect people who don't use it |
Per maintainer request on PR apache#22689, the `LambdaExpr::try_new` / `expressions::lambda(...)` signature change (adding `outer_columns_count`) is being reviewed separately in apache#22853 because it's a breaking change to the physical-expr public API and warrants its own attention, distinct from this additive UDF. This commit reverts the `lambda.rs` / `higher_order_function.rs` / `planner.rs` files to their `upstream/main` state, removes the sqllogictest file (every query in it uses `(k, v) -> body` lambdas that require the upstream fix), and marks the unit tests that exercise multi-parameter lambdas with `#[ignore = "blocked on apache#22853: multi-param lambda projection fix"]`. `transform_values_uses_keys_via_case` and `transform_values_all_null_rows_returns_null_array` still pass because the former references both `k` and `v` (so projection is a no-op) and the latter short-circuits before evaluating the lambda. This PR will be rebased onto main once apache#22853 merges, at which point the ignore markers will be removed and the sqllogictest file restored.
|
I get the feeling that this UDF might be too specific to be contributed and maintained in this repo. Before moving forward, is there any precedence you see in other engines for this function? typically this project aims to host functions present in engines like Postgres, or Spark. |
Rationale for this change
We want to give SQL users a way to rewrite the values of a map without re-implementing the key/value plumbing themselves. DataFusion already has
array_transformfrom #18921 as a precedent for higher-order array functions, so this is the natural map-side equivalent.What changes are included in this PR?
transform_values(map, (k, v) -> expr)as a higher-order UDF indatafusion-functions-nested. The function applies the lambda to every entry, returns a new map with the original keys and the lambda's results as the new values, and propagates null rows.get_map_key_value_fieldsinfunctions-nested/src/utils.rsso other map UDFs can share the "get the key/value field refs out of aMap" pattern.The breaking change to
LambdaExprthat originally lived in this PR has been extracted to #22853.